import matplotlib.pyplot as plt
import torch
from torch import nn
from torchvision import transforms as T
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl
from torchmetrics import Accuracy



class DataM(pl.LightningDataModule):
    def __init__(self, data_dir: str,
                 batch_size: int,
                 num_workers: int ):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        transform = T.Compose([T.ToTensor()])
        self.ds_test = MNIST(self.data_dir, train=False, transform=transform, download=True)
        self.ds_predict = MNIST(self.data_dir, train=False, transform=transform, download=True)
        ds_full = MNIST(self.data_dir, train=True, transform=transform, download=True)
        # _,ds_full = random_split(ds_full, [55000, 5000])
        # self.ds_train, self.ds_val = random_split(ds_full, [55000, 5000])
        self.ds_train, self.ds_val = ds_full,ds_full

    def train_dataloader(self):
        return DataLoader(self.ds_train, batch_size=self.batch_size,
                          shuffle=True, num_workers=self.num_workers,
                          pin_memory=True)

    def val_dataloader(self):
        return DataLoader(self.ds_val, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True)

    def test_dataloader(self):
        return DataLoader(self.ds_test, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True)

    def predict_dataloader(self):
        return DataLoader(self.ds_predict, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True)

if __name__ == '__main__':

    data_mnist = DataM("../../data/", 32, 0)
    data_mnist.setup()
    images, labels = next(iter(data_mnist.train_dataloader()))
    print(images.shape)


